package com.allenkerr.rpc.registry;

import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadLocalRandom;

/**
 * 功能：client 递归设置对server的节点监听，获取server的地址
 * client 监听 server 节点的变化，可以实现负载均衡
 * <p>
 * Created by 小柯(Allen Kerr) on 17.5.23 18:56.
 */
public class ServiceDiscovery {

    private static final Logger LOGGER = LoggerFactory.getLogger(ServiceDiscovery.class);

    private CountDownLatch latch = new CountDownLatch(1);

    private volatile List<String> dataList = new ArrayList<>();

    private String registryAddress;

    /**
     * 构造函数，参数：zookeeper注册地址
     *
     * @param registryAddress
     */
    public ServiceDiscovery(String registryAddress) {
        this.registryAddress = registryAddress;
        //连接到zookeeper
        ZooKeeper zooKeeper = connectServer();
        //如果连接成功，递归设置对server的节点监听，获取server的地址
        if (zooKeeper != null) {
            watchNode(zooKeeper);
        }
    }

    /**
     * 发现新节点
     *
     * @return data
     */
    public String discover() {
        String data = null;
        int size = dataList.size();
        //从成员变量dataList中获取data
        if (size > 0) {
            if (size == 1) {
                data = dataList.get(0);
                LOGGER.debug("Using only data: {}", data);
            } else {
                data = dataList.get(ThreadLocalRandom.current().nextInt(size));
                LOGGER.debug("Using random data: {}", data);
            }
        }
        return data;
    }

    /**
     * 创建zookeeper的连接
     *
     * @return
     */
    private ZooKeeper connectServer() {
        ZooKeeper zooKeeper = null;
        try {
            zooKeeper = new ZooKeeper(registryAddress, RpcConstant.ZK_SESSION_TIMEOUT,
                    new Watcher() {
                        @Override
                        public void process(WatchedEvent watchedEvent) {
                            if (watchedEvent.getState() == Event.KeeperState.SyncConnected) {
                                latch.countDown();
                            }
                        }
                    });
            latch.await();
        } catch (Exception e) {
            LOGGER.error("", e);
        }
        return zooKeeper;
    }

    /**
     * 监听zookeeper
     *
     * @param zooKeeper
     */
    private void watchNode(ZooKeeper zooKeeper) {
        try {
            //获取所有子节点
            List<String> nodeList = zooKeeper.getChildren(RpcConstant.ZK_REGISTRY_PATH, new Watcher() {
                @Override
                public void process(WatchedEvent watchedEvent) {
                    //监控节点变化
                    if (watchedEvent.getType() == Event.EventType.NodeChildrenChanged)
                        watchNode(zooKeeper);
                }
            });
            List<String> dataList = new ArrayList<>();
            //循环子节点
            for (String node : nodeList) {
                //获取节点中的服务器地址
                byte[] data = zooKeeper.getData(RpcConstant.ZK_REGISTRY_PATH + "/" + node, false, null);
                //添加到list中
                dataList.add(new String(data));
            }
            LOGGER.debug("Node Data: {}", dataList);
            //将节点信息记录到成员变量中
            this.dataList = dataList;
        } catch (Exception e) {
            LOGGER.error("", e);
        }


    }

}
